#ifndef __TCP_SERVER_H__
#define __TCP_SERVER_H__

#include <iostream>
#include <string>
#include <cstring>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <vector>

using namespace std;

class TCPServer {
private:
    int port;                // 服务器监听的端口号
    int server_fd;           // 服务器套接字文件描述符
    vector<int> client_fds;  // 客户端套接字文件描述符列表
    struct sockaddr_in server_addr;  // 服务器地址结构

public:
    // 构造函数：初始化端口号和套接字
    TCPServer(int port) : port(port), server_fd(-1) {}

    // 析构函数：关闭所有客户端连接并关闭服务器套接字
    ~TCPServer() {
        close_connections();
    }

    // 启动服务器：创建套接字，绑定并开始监听
    bool start() {
        server_fd = socket(AF_INET, SOCK_STREAM, 0);
        if (server_fd == -1) {
            cerr << "创建套接字失败！" << endl;
            return false;
        }

        server_addr.sin_family = AF_INET;
        server_addr.sin_addr.s_addr = INADDR_ANY;
        server_addr.sin_port = htons(port);

        if (bind(server_fd, (struct sockaddr*)&server_addr, sizeof(server_addr)) == -1) {
            cerr << "绑定失败！" << endl;
            return false;
        }

        if (listen(server_fd, 5) == -1) {
            cerr << "监听失败！" << endl;
            return false;
        }

        cout << "服务器已启动，等待客户端连接..." << endl;
        return true;
    }

    // 接受客户端连接
    bool acceptClient() {
        struct sockaddr_in client_addr;
        socklen_t client_len = sizeof(client_addr);
        int client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &client_len);
        if (client_fd == -1) {
            cerr << "接受客户端连接失败！" << endl;
            return false;
        }

        client_fds.push_back(client_fd);
        cout << "客户端连接成功！" << endl;
        return true;
    }

    // 接收客户端发送的数据
    bool receiveData(string &data) {
        char buffer[1024];
        ssize_t bytes_received = recv(client_fds.back(), buffer, sizeof(buffer) - 1, 0);
        if (bytes_received > 0) {
            buffer[bytes_received] = '\0';
            data = buffer;
            return true;
        }
        return false;
    }

    // 向客户端发送数据
    bool sendData(const string &data) {
        ssize_t bytes_sent = send(client_fds.back(), data.c_str(), data.length(), 0);
        return bytes_sent > 0;
    }

    // 重载 << 操作符，向客户端发送数据
    TCPServer& operator<<(const string &data) {
        sendData(data);
        return *this;
    }

    // 重载 >> 操作符，从客户端接收数据
    TCPServer& operator>>(string &data) {
        receiveData(data);
        return *this;
    }

    // 关闭所有连接
    void close_connections() {
        for (int client_fd : client_fds) {
            close(client_fd);
        }
        close(server_fd);
        cout << "服务器已关闭。" << endl;
    }
};

#endif // TCP_SERVER_H
